import gym
import numpy as np


env = gym.make('FrozenLake-v1', desc=None, map_name="4x4", is_slippery=True)
observation = env.observation_space
action = env.action_space
print(observation)
print(action)

eps = 1
policy = np.ones(16)
value = np.zeros(16)

'''
for state in range(0,16):
    for action in range(0,4):
        for state_probability, next_state, reward, terminated in env.P[state][action]:
            print(state, action)
            print(state_probability, next_state, reward, terminated)
'''

while eps > 0.001:
    for j in range(0, 10):
        new_value = np.copy(value)
        for state in range(0, 16):
            value[state] = sum([state_probability*(reward+new_value[next_state])
                for state_probability, next_state, reward, terminated in env.P[state][policy[state]]])

    print(value)

    policy_k = np.array(policy)
    print(policy_k)

    for state in range(0, 16):
        Q_table = np.zeros(4)
        for action in range(0, 4):
            for state_probability, next_state, reward, terminated  in env.P[state][action]:
                Q_table[action] += (state_probability *(reward+value[next_state]))
        policy[state] = np.argmax(Q_table)
    print(policy)
    eps = np.linalg.norm(policy - policy_k)